import os
import cv2
import numpy as np
from facenet_pytorch import MTCNN, InceptionResnetV1
import pickle
import torch.nn as nn
from PIL import Image, ImageDraw, ImageFont
import streamlit as st
import torch
import time

# 配置页面
st.set_page_config(page_title="人脸识别系统", layout="wide")

# 确保中文显示正常
st.title("人脸识别系统")
st.write("实时检测和识别人脸")

# 缓存资源，避免重复加载
@st.cache_resource
def load_models():
    try:
        mtcnn = MTCNN()
        # 初始化ResNet模型
        resnet = InceptionResnetV1(pretrained="vggface2").eval().to("cpu")
        return mtcnn, resnet
    except Exception as e:
        st.error(f"模型加载失败: {e}")
        return None, None

mtcnn, resnet = load_models()
if mtcnn is None or resnet is None:
    st.stop()

# 加载人脸数据库
@st.cache_data
def load_face_database(data_path="static/images"):
    face_dict = {}
    if not os.path.exists(data_path):
        st.warning(f"人脸数据库路径不存在: {data_path}")
        return face_dict
    try:
        for entry in os.scandir(data_path):
            face_name = os.path.splitext(entry.name)[0]
            with open(entry.path, "rb") as f:
                face_tensor = pickle.load(f)
            face_dict[face_name] = face_tensor
        return face_dict
    except Exception as e:
        st.error(f"加载人脸数据库失败: {e}")
        return face_dict

face_dict = load_face_database()

# 与人脸数据库进行匹配
def pk_face(frame, threshold=0.6):
    global face_dict
    if len(face_dict) == 0:
        return ["未知"]

    # 转换BGR到RGB
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 检测人脸
    boxes, probs, points = mtcnn.detect(frame_rgb, landmarks=True)
    result = []

    if boxes is not None:
        for i in range(len(boxes)):
            box = boxes[i].reshape(1, 4)
            try:
                # 提取人脸特征
                face_tensor = mtcnn.extract(frame_rgb, box, save_path=None)
                if face_tensor is None:
                    result.append("检测失败")
                    continue
                # 计算特征嵌入
                face_embed = resnet(face_tensor.unsqueeze(0).to("cpu"))

                max_score = float("-inf")
                max_score_idx = -1

                # 对比数据库
                for j, ds_face in enumerate(face_dict.values()):
                    score = nn.functional.cosine_similarity(face_embed, ds_face, dim=1).item()
                    if score > max_score:
                        max_score = score
                        max_score_idx = j

                # 判断结果
                if max_score > threshold:
                    result.append(f"{list(face_dict.keys())[max_score_idx]}({max_score:.2f})")
                else:
                    result.append(f"未知({max_score:.2f})")
            except Exception as e:
                result.append(f"处理错误: {str(e)[:10]}")

    return result

# 字体配置
font = ImageFont.truetype("simkai.ttf", 28)

# 在图像上绘制结果
def draw_results(frame, boxes, face_names):
    if font is None:
        return frame

    img_pil = Image.fromarray(frame)
    draw = ImageDraw.Draw(img_pil)

    if boxes is not None:
        for i, box in enumerate(boxes):
            x1, y1, x2, y2 = box.astype(int)
            # 绘制人脸框
            draw.rectangle([x1, y1, x2, y2], outline=(0, 0, 255), width=2)
            # 绘制姓名
            if i < len(face_names):
                draw.text((x1, y1 - 30), face_names[i], fill=(255, 0, 0), font=font)
    return np.array(img_pil, dtype=np.uint8)

# 主函数
def main():
    # 侧边栏设置
    st.sidebar.header("设置")
    threshold = st.sidebar.slider("识别阈值", 0.1, 1.0, 0.6, 0.05)

    # 在侧边栏添加名字输入框
    text = st.sidebar.text_input("请输入名字", key="face_name")

    # 主界面按钮布局
    col1, col2 = st.columns(2)
    with col1:
        st.subheader("操作按钮")
        start_detection = st.button("开始检测")
        get_detection = st.button("捕获人脸")
        save_detection = st.button("保存")

    with col2:
        st.subheader("摄像头控制")
        start_camera = st.button("开始识别")
        stop_camera = st.button("停止识别")

    # 初始化状态
    if 'running' not in st.session_state:
        st.session_state.running = False
    if 'captured_face' not in st.session_state:
        st.session_state.captured_face = None
    if 'face_embed' not in st.session_state:
        st.session_state.face_embed = None
    if 'capture_success' not in st.session_state:
        st.session_state.capture_success = False
    if 'start_time' not in st.session_state:
        st.session_state.start_time = time.time()

    # 控制摄像头状态
    if start_camera or start_detection:
        st.session_state.running = True
        st.success("摄像头开启中")
    if stop_camera:
        st.session_state.running = False

    # 摄像头显示区域
    col3, col4 = st.columns(2)
    with col3:
        st.subheader("摄像头画面")
        camera_placeholder = st.empty()
    with col4:
        st.subheader("识别结果")
        result_placeholder = st.empty()

    cap = None
    try:
        if st.session_state.running:
            # 打开摄像头
            cap = cv2.VideoCapture(0)
            if not cap.isOpened():
                st.error("无法打开摄像头")
                return

            st.session_state.frame_count = 0
            st.session_state.start_time = time.time()

            while st.session_state.running:
                ret, frame = cap.read()
                if not ret:
                    st.error("无法获取摄像头画面")
                    break

                frame = cv2.flip(frame, 1)
                start_time = time.time()

                # 人脸识别
                face_names = pk_face(frame, threshold)

                # 绘制结果
                boxes, probs, points = mtcnn.detect(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), landmarks=True)
                result_frame = draw_results(frame, boxes, face_names)

                # 捕获人脸
                if get_detection and not st.session_state.capture_success:
                    if boxes is not None and len(boxes) > 0:
                        idx = probs.argmax()
                        try:
                            # 提取人脸特征
                            face_tensor = mtcnn.extract(
                                cv2.cvtColor(frame, cv2.COLOR_BGR2RGB),
                                boxes[idx].reshape(1, -1),
                                save_path=None
                            )
                            if face_tensor is not None:
                                # 计算特征嵌入
                                face_embed = resnet(face_tensor.unsqueeze(0).to("cpu"))

                                # 保存状态
                                st.session_state.captured_face = frame.copy()
                                st.session_state.face_embed = face_embed
                                st.session_state.capture_success = True
                                st.success("人脸捕获成功，请输入名字并点击保存")
                            else:
                                st.error("未检测到有效人脸，请重试")
                        except Exception as e:
                            st.error(f"捕获人脸失败: {e}")
                    else:
                        st.error("未检测到人脸，请重试")

                # 保存人脸逻辑
                if save_detection and st.session_state.capture_success:
                    # 从会话状态获取名字
                    text = st.session_state.face_name
                    if not text:
                        st.error("请先输入名字")
                    else:
                        save_dir = "static/images"
                        if not os.path.exists(save_dir):
                            os.makedirs(save_dir)

                        try:
                            # 保存人脸特征
                            with open(f"{save_dir}/{text}.embed", "wb") as f:
                                pickle.dump(st.session_state.face_embed.cpu().detach(), f)

                            global face_dict
                            face_dict = load_face_database()

                            st.success(f"成功保存 {text} 的人脸数据")
                            st.session_state.captured_face = None
                            st.session_state.face_embed = None
                            st.session_state.capture_success = False
                        except Exception as e:
                            st.error(f"保存失败: {e}")

                # 显示图像
                camera_placeholder.image(frame, channels="BGR", use_container_width=True)
                result_placeholder.image(result_frame, channels="BGR", use_container_width=True)

                # 帧率
                process_time = time.time() - start_time
                if process_time < 0.03:  # 30FPS
                    time.sleep(0.03 - process_time)

    except Exception as e:
        st.error(f"程序运行出错: {e}")

    finally:
        # 释放资源
        if cap is not None and cap.isOpened():
            cap.release()
        st.session_state.running = False
        st.success("程序已停止")

if __name__ == "__main__":
    main()